// Copyright (c) 2012 The Chromium Authors. All rights reserved.
// Use of this source code is governed by a BSD-style license that can be
// found in the LICENSE file.

#include "src/base64.h"
#include <algorithm>
#include "src/modp_b64.h"

void Base64Encode(const std::string &input, std::string *output) {
    std::string temp;
    temp.resize(modp_b64_encode_len(input.size()));  // makes room for null byte

    // modp_b64_encode_len() returns at least 1, so temp[0] is safe to use.
    size_t output_size = modp_b64_encode(&(temp[0]), input.data(), input.size());

    temp.resize(output_size);  // strips off null byte
    output->swap(temp);
}

bool Base64Decode(const std::string &input, std::string *output) {
    std::string temp;
    temp.resize(modp_b64_decode_len(input.size()));

    // does not null terminate result since result is binary data!
    size_t input_size = input.size();
    size_t output_size = modp_b64_decode(&(temp[0]), input.data(), input_size);
    if (output_size == MODP_B64_ERROR) return false;

    temp.resize(output_size);
    output->swap(temp);
    return true;
}

const char kPaddingChar = '=';

// Base64url maps {+, /} to {-, _} in order for the encoded content to be safe
// to use in a URL. These characters will be translated by this implementation.
const char kBase64Chars[] = "+/";
const char kBase64UrlSafeChars[] = "-_";

void ReplaceChars(const std::string &input, char old_char, char new_char, std::string *output) {
    std::transform(input.begin(), input.end(), output->begin(),
                   [old_char, new_char](char cc) -> char {
                       if (cc == old_char) {
                           return new_char;
                       }
                       return cc;
                   });
}

void Base64UrlEncode(const std::string &input, Base64UrlEncodePolicy policy, std::string *output) {
    Base64Encode(input, output);

    ReplaceChars(*output, '+', '-', output);
    ReplaceChars(*output, '/', '_', output);

    switch (policy) {
    case Base64UrlEncodePolicy::INCLUDE_PADDING:
        // The padding included in |*output| will not be amended.
        break;
    case Base64UrlEncodePolicy::OMIT_PADDING:
        // The padding included in |*output| will be removed.
        const size_t last_non_padding_pos = output->find_last_not_of(kPaddingChar);
        if (last_non_padding_pos != std::string::npos) output->resize(last_non_padding_pos + 1);

        break;
    }
}

bool Base64UrlDecode(const std::string &input, Base64UrlDecodePolicy policy, std::string *output) {
    // Characters outside of the base64url alphabet are disallowed, which includes
    // the {+, /} characters found in the conventional base64 alphabet.
    if (input.find_first_of(kBase64Chars) != std::string::npos) return false;

    const size_t required_padding_characters = input.size() % 4;
    const bool needs_replacement = input.find_first_of(kBase64UrlSafeChars) != std::string::npos;

    switch (policy) {
    case Base64UrlDecodePolicy::REQUIRE_PADDING:
        // Fail if the required padding is not included in |input|.
        if (required_padding_characters > 0) return false;
        break;
    case Base64UrlDecodePolicy::IGNORE_PADDING:
        // Missing padding will be silently appended.
        break;
    case Base64UrlDecodePolicy::DISALLOW_PADDING:
        // Fail if padding characters are included in |input|.
        if (input.find_first_of(kPaddingChar) != std::string::npos) return false;
        break;
    }

    // If the string either needs replacement of URL-safe characters to normal
    // base64 ones, or additional padding, a copy of |input| needs to be made in
    // order to make these adjustments without side effects.
    if (required_padding_characters > 0 || needs_replacement) {
        std::string base64_input;

        std::size_t base64_input_size = input.size();
        if (required_padding_characters > 0) base64_input_size += 4 - required_padding_characters;

        base64_input.reserve(base64_input_size);
        base64_input.append(input.data(), input.size());

        // Substitute the base64url URL-safe characters to their base64 equivalents.
        ReplaceChars(base64_input, '-', '+', &base64_input);
        ReplaceChars(base64_input, '_', '/', &base64_input);

        // Append the necessary padding characters.
        base64_input.resize(base64_input_size, '=');

        return Base64Decode(base64_input, output);
    }

    return Base64Decode(input, output);
}
